#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Create by zhang
# Create on 2022/8/9 12:01
from typing import List, Tuple

import pandas as pd
from pandas import DataFrame
from numpy import nan, array, ndarray

from core.dataClasses import StockTradeDataColumnName
from core.index_base import IndexBase

SMA_COLUMN_PREFIX = 'sma'

# TODO 移动平均的移动平均

class SMA(IndexBase):
    """
    Simple Moving Average 简单移动平均数：将一组股价值相加，再除以股价的个数
    \[SMA_{t = 5} = \frac{{{p_1} + {p_2} + {p_3} + {p_4} + {p_5}}}{5}\]
    """

    def __init__(self, period=5, key: str = StockTradeDataColumnName.CLOSE):
        super(SMA, self).__init__()
        self.period = period
        self.key = key

    def compute(self, data: DataFrame):
        sma_key = f"{self.key}_sma_{self.period}"
        self.added_keys.append(sma_key)
        data[sma_key] = nan
        data[sma_key] = data[self.key].rolling(self.period).mean()


class WMA(IndexBase):
    """
    Weighted Moving Average 加权移动平均数：其与简单移动平均数不同，其并不是简单地把股票价格相加，而是对股价赋予一定的权重再相加
    离当前时间越近的数据越具有代表性，越久远的数据越没有代表性。因此，再预测股价时，不同时期的股价数量具有不同的代表性。为了表示其代表性的高低，可以考虑对股价赋予一定的权重，再求平均值。
    \[WMA_{t = 5} = {w_1}{p_1} + {w_2}{p_2} + {w_3}{p_3} + {w_4}{p_4} + {w_5}{p_5}\]，式中，w_i为股价数据的权重且\[{w_1} + {w_2} + {w_3} + {w_4} + {w_5} = 1\]
    """

    def __init__(self, weights: Tuple[int] = (1, 2, 3, 4, 5,), key: str = StockTradeDataColumnName.CLOSE):
        """
        :param weights: 为int列表，其值为其对应位置所占的权重，列表数量为计算周期
        :param key: 为计算的键名
        """
        arr = array(weights)
        weights: ndarray = arr / sum(arr)
        weights = weights.tolist()
        self.weights = weights
        self.key = key

    def compute(self, data: DataFrame):
        wma_key = 'wma' + str(len(self.weights))
        self.added_keys.append(wma_key)
        data[wma_key] = nan
        index_list = data.index.values
        for i in range(0, len(index_list)):
            sum_val = 0
            for j in range(1, len(self.weights)):
                if i - j >= 0:
                    sum_val += data.loc[index_list[i - j]][self.key] * self.weights[i - j]
            sum_val += data.loc[index_list[i]][self.key] * self.weights[-1]
            data.loc[index_list[i], wma_key] = sum_val


class EWMA(IndexBase):
    """
    Exponential Weighted Moving Average 指数加权移动平均数：指数加权移动平均数相当于一种特别的加权移动平均。在此以5日指数移动平均为例。
    我们先给定一个权重值，比如0.2。由加权移动平均的定义可知，我们无法得到前4期的5日加权移动平均数，而只能得到从5期开始的平均数。
    假设p_t表示股票第t期的价格，我们从第k期开始计算股价的加权移动平均数，且第k期的平均数计算非常简单，比如用前k期的股价之简单平均数求得，即：
    \[EWMA_{t = k} = \frac{{{p_1} + {p_2} + {p_3} + {p_4} + {p_k}}}{k}\]
    从第k+1期开始，每一期的移动平均数为当前股价与上一期移动平均数之加权平均，权重分别为\[\alpha \]和\[1 - \alpha \]，即：
    \[{\rm{EWMA}}_{t = k + 1}  =  {{\rm{p}}_{k + 1}}\alpha   +  {\rm EWMA}_{t = k}(1 - \alpha )\]
    第k+2期指数加权移动平均数为：
    \[{\rm EWMA}_{t = k + 2}  =  {{\rm{p}}_{k + 2}}\alpha   +  {\rm {EWMA}}_{{\rm{t}} = k + 1}}(1 - \alpha )\]
    """

    def __init__(self, period=5, alpha=0.8, key: str = StockTradeDataColumnName.CLOSE):
        """

        :param period: 计算周期
        :param alpha: 周期最后一条数据的所占的权重
        :param key: 基准键名
        """
        super(EWMA, self).__init__()
        self.period = period
        self.key = key
        self.alpha = alpha

    def compute(self, data: DataFrame):
        # TODO 算法完善
        sma_key = 'ewma' + str(self.period)
        self.added_keys.append(sma_key)
        data[sma_key] = nan
        index_list = data.index.values
        index_list = [str(i)[:10] for i in index_list]
        sum_val = 0
        for i in range(0, self.period):
            sum_val += data.loc[index_list[i]][self.key]
        data.loc[index_list[self.period - 1], sma_key] = sum_val / self.period
        data.loc[index_list[self.period]:, sma_key] = data.loc[index_list[self.period]:, self.key]

        data.loc[index_list[self.period - 1]:, sma_key] = data.loc[index_list[self.period - 1]:, sma_key].ewm(
            alpha=self.alpha, adjust=False).mean()
        print()


class MACD(IndexBase):
    """

    """
    def __init__(self):
        pass

    def compute(self, data:DataFrame):
        pass